import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from datetime import datetime, timedelta
import warnings

warnings.filterwarnings('ignore')
plt.style.use('seaborn')
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

# 设置Tushare token (需要先在Tushare官网注册获取)
ts.set_token('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')
pro = ts.pro_api()


class StockStrategy:
    def __init__(self):
        self.data = {}
        self.results = {}

    def get_data(self, stock_codes, start_date, end_date):
        """
        获取股票日线数据
        :param stock_codes: 股票代码列表，如 ['600519.SH', '000858.SZ']
        :param start_date: 开始日期，格式 'YYYYMMDD'
        :param end_date: 结束日期，格式 'YYYYMMDD'
        """
        print("正在获取股票数据...")

        for code in stock_codes:
            try:
                df = pro.daily(ts_code=code, start_date=start_date, end_date=end_date)
                df = df.sort_values('trade_date')
                df['trade_date'] = pd.to_datetime(df['trade_date'])
                df.set_index('trade_date', inplace=True)
                df['pct_chg'] = df['close'].pct_change() * 100
                self.data[code] = df
                print(f"成功获取 {code} 数据，共 {len(df)} 条记录")
            except Exception as e:
                print(f"获取 {code} 数据失败: {e}")

    def process_data(self, short_window=5, long_window=20):
        """
        数据处理，计算技术指标
        :param short_window: 短期均线窗口
        :param long_window: 长期均线窗口
        """
        print("\n正在处理数据，计算技术指标...")

        for code, df in self.data.items():
            # 计算均线
            df['short_ma'] = df['close'].rolling(window=short_window).mean()
            df['long_ma'] = df['close'].rolling(window=long_window).mean()

            # 计算均线交叉信号
            df['signal'] = 0
            df['signal'][short_window:] = np.where(
                df['short_ma'][short_window:] > df['long_ma'][short_window:], 1, 0)

            # 生成交易信号 (1: 买入, -1: 卖出, 0: 持有)
            df['trade_signal'] = df['signal'].diff()

            # 标记买卖点
            df['buy_signal'] = np.where(df['trade_signal'] == 1, df['close'], np.nan)
            df['sell_signal'] = np.where(df['trade_signal'] == -1, df['close'], np.nan)

            print(f"{code} 数据处理完成")

    def backtest(self, initial_capital=100000, stop_loss=0.1, stop_gain=0.2):
        """
        策略回测
        :param initial_capital: 初始资金
        :param stop_loss: 止损比例
        :param stop_gain: 止盈比例
        """
        print("\n开始回测...")

        for code, df in self.data.items():
            print(f"\n回测股票: {code}")

            # 初始化回测变量
            capital = initial_capital
            position = 0
            buy_price = 0
            trade_records = []
            equity_curve = []

            for i, row in df.iterrows():
                # 记录每日权益
                if position > 0:
                    current_value = position * row['close']
                    equity_curve.append(current_value + (capital - position * buy_price))
                else:
                    equity_curve.append(capital)

                # 检查止损止盈
                if position > 0:
                    current_return = (row['close'] - buy_price) / buy_price
                    if current_return <= -stop_loss or current_return >= stop_gain:
                        # 卖出止损/止盈
                        capital = position * row['close']
                        position = 0
                        buy_price = 0
                        trade_records.append({
                            'date': i,
                            'type': 'sell',
                            'price': row['close'],
                            'reason': 'stop_loss' if current_return <= -stop_loss else 'stop_gain',
                            'capital': capital
                        })
                        continue

                # 执行交易信号
                if row['trade_signal'] == 1 and capital > 0:  # 买入信号
                    buy_price = row['close']
                    position = capital / buy_price
                    capital = 0
                    trade_records.append({
                        'date': i,
                        'type': 'buy',
                        'price': buy_price,
                        'reason': 'signal',
                        'capital': position * buy_price
                    })
                elif row['trade_signal'] == -1 and position > 0:  # 卖出信号
                    capital = position * row['close']
                    position = 0
                    buy_price = 0
                    trade_records.append({
                        'date': i,
                        'type': 'sell',
                        'price': row['close'],
                        'reason': 'signal',
                        'capital': capital
                    })

            # 最后一天平仓
            if position > 0:
                capital = position * df.iloc[-1]['close']
                trade_records.append({
                    'date': df.index[-1],
                    'type': 'sell',
                    'price': df.iloc[-1]['close'],
                    'reason': 'close',
                    'capital': capital
                })

            # 计算回测结果
            final_value = capital
            total_return = (final_value - initial_capital) / initial_capital * 100
            trade_count = len([t for t in trade_records if t['type'] == 'buy'])
            win_trades = len([t for t in trade_records
                              if t['type'] == 'sell' and t['reason'] in ['signal', 'stop_gain']
                              and t['price'] > [b['price'] for b in trade_records
                                                if b['type'] == 'buy' and b['date'] < t['date']][-1]])

            win_rate = win_trades / trade_count * 100 if trade_count > 0 else 0

            # 保存结果
            self.results[code] = {
                'initial_capital': initial_capital,
                'final_value': final_value,
                'total_return': total_return,
                'trade_count': trade_count,
                'win_rate': win_rate,
                'trade_records': trade_records,
                'equity_curve': equity_curve,
                'dates': df.index
            }

            print(f"初始资金: {initial_capital:.2f}")
            print(f"最终价值: {final_value:.2f}")
            print(f"总收益率: {total_return:.2f}%")
            print(f"交易次数: {trade_count}")
            print(f"胜率: {win_rate:.2f}%")

    def visualize_signals(self, stock_code):
        """
        可视化交易信号
        :param stock_code: 股票代码
        """
        if stock_code not in self.data:
            print(f"没有找到 {stock_code} 的数据")
            return

        df = self.data[stock_code]
        plt.figure(figsize=(14, 8))

        # 绘制K线
        plt.plot(df.index, df['close'], label='收盘价', color='black', alpha=0.7)
        plt.plot(df.index, df['short_ma'], label=f'短期均线({df["short_ma"].notna().idxmax()})', color='blue',
                 alpha=0.5)
        plt.plot(df.index, df['long_ma'], label=f'长期均线({df["long_ma"].notna().idxmax()})', color='red', alpha=0.5)

        # 标记买卖点
        plt.scatter(df.index, df['buy_signal'], label='买入信号', marker='^', color='green', alpha=1)
        plt.scatter(df.index, df['sell_signal'], label='卖出信号', marker='v', color='red', alpha=1)

        plt.title(f'{stock_code} 双均线策略交易信号')
        plt.xlabel('日期')
        plt.ylabel('价格')
        plt.legend()
        plt.grid(True)
        plt.show()

    def visualize_equity_curve(self, stock_codes=None):
        """
        可视化资金曲线
        :param stock_codes: 要显示的股票代码列表，None表示显示所有
        """
        if stock_codes is None:
            stock_codes = list(self.results.keys())

        plt.figure(figsize=(14, 7))

        for code in stock_codes:
            if code in self.results:
                result = self.results[code]
                plt.plot(result['dates'], result['equity_curve'],
                         label=f'{code} (收益率:{result["total_return"]:.2f}%)')

        plt.title('策略资金曲线对比')
        plt.xlabel('日期')
        plt.ylabel('资金')
        plt.legend()
        plt.grid(True)
        plt.show()

    def report(self):
        """
        生成回测报告
        """
        print("\n回测报告:")
        print("-" * 80)
        print(f"{'股票代码':<10} {'初始资金':>10} {'最终价值':>10} {'总收益率(%)':>12} {'交易次数':>8} {'胜率(%)':>8}")
        print("-" * 80)

        for code, result in self.results.items():
            print(f"{code:<10} {result['initial_capital']:>10.2f} {result['final_value']:>10.2f} "
                  f"{result['total_return']:>12.2f} {result['trade_count']:>8} {result['win_rate']:>8.2f}")

        # 计算基准收益率
        print("\n基准收益率(买入持有策略):")
        print("-" * 80)
        print(f"{'股票代码':<10} {'初始价格':>10} {'最终价格':>10} {'收益率(%)':>12}")
        print("-" * 80)

        for code, df in self.data.items():
            initial_price = df.iloc[0]['close']
            final_price = df.iloc[-1]['close']
            return_pct = (final_price - initial_price) / initial_price * 100
            print(f"{code:<10} {initial_price:>10.2f} {final_price:>10.2f} {return_pct:>12.2f}")


# 示例使用
if __name__ == "__main__":
    # 设置回测参数
    stock_codes = ['600519.SH', '000858.SZ', '600036.SH']  # 贵州茅台, 五粮液, 招商银行
    start_date = '20240101'
    end_date = '20250113'

    # 创建策略实例
    strategy = StockStrategy()

    # 获取数据
    strategy.get_data(stock_codes, start_date, end_date)

    # 处理数据，计算指标
    strategy.process_data(short_window=10, long_window=30)

    # 运行回测
    strategy.backtest(initial_capital=100000, stop_loss=0.08, stop_gain=0.15)

    # 可视化
    for code in stock_codes:
        strategy.visualize_signals(code)

    strategy.visualize_equity_curve()

    # 生成报告
    strategy.report()